from code_generator import CodeGeneratorAI

# 准备训练数据（这里是简单示例，实际应用需更多数据）
training_data = [
    # 基础运算
    "def add_numbers(a, b):\n    # 计算两个数的和\n    return a + b",
    "def multiply_numbers(a, b):\n    # 计算两个数的乘积\n    return a * b",
    "def is_even(n):\n    # 判断一个数是否为偶数\n    return n % 2 == 0",
    "def square(n):\n    # 计算一个数的平方\n    return n ** 2",
    "def factorial(n):\n    # 计算一个数的阶乘（递归实现）\n    if n == 0:\n        return 1\n    else:\n        return n * factorial(n-1)",

    # 模块导入
    "import math\n\ndef calculate_circle_area(radius):\n    # 计算圆的面积，使用math模块的pi常量\n    return math.pi * radius ** 2",
    "import random\n\ndef roll_dice(sides=6):\n    # 模拟掷骰子，返回1到sides之间的随机整数\n    return random.randint(1, sides)",
    "import os\n\ndef create_directory(path):\n    # 创建目录，如果目录不存在\n    if not os.path.exists(path):\n        os.makedirs(path)",
    "import numpy as np\n\ndef calculate_mean(numbers):\n    # 计算列表的平均值，使用numpy库\n    return np.mean(numbers)",
    "import pandas as pd\n\ndef read_csv_file(file_path):\n    # 读取CSV文件并返回DataFrame\n    return pd.read_csv(file_path)",

    # Pygame游戏开发
    "def init_game(title, width, height):\n    # 初始化Pygame并创建游戏窗口\n    import pygame\n    pygame.init()\n    screen = pygame.display.set_mode((width, height))\n    pygame.display.set_caption(title)\n    clock = pygame.time.Clock()\n    return screen, clock",
    "def load_image(path, scale=None):\n    # 加载并缩放图像\n    import pygame\n    image = pygame.image.load(path).convert_alpha()\n    if scale:\n        image = pygame.transform.scale(image, scale)\n    return image",
    "class GameObject(pygame.sprite.Sprite):\n    # 游戏对象基类，继承自Pygame的Sprite类\n    def __init__(self, x, y, width, height, color=(255, 255, 255)):\n        super().__init__()\n        self.image = pygame.Surface((width, height))\n        self.image.fill(color)\n        self.rect = self.image.get_rect()\n        self.rect.x = x\n        self.rect.y = y\n    def update(self):\n        # 更新对象状态的方法（需在子类中实现）\n        pass",
    "def check_collision(sprite1, sprite2):\n    # 检查两个精灵是否发生碰撞\n    return pygame.sprite.collide_rect(sprite1, sprite2)",
    "def animate_sprite(sprite, frames, frame_rate, loop=True):\n    # 实现精灵动画效果\n    if not hasattr(sprite, 'animation_frame'):\n        sprite.animation_frame = 0\n        sprite.animation_timer = 0\n    sprite.animation_timer += 1\n    if sprite.animation_timer >= frame_rate:\n        sprite.animation_frame = (sprite.animation_frame + 1) % len(frames)\n        sprite.animation_timer = 0\n        sprite.image = frames[sprite.animation_frame]\n    if not loop and sprite.animation_frame == len(frames) - 1:\n        return False\n    return True",

    # 数据处理
    "import pandas as pd\n\ndef analyze_sales_data(file_path):\n    # 分析销售数据，计算总销售额和平均销售额\n    df = pd.read_csv(file_path)\n    total_sales = df['sales'].sum()\n    avg_sales = df['sales'].mean()\n    return {'total': total_sales, 'average': avg_sales}",
    "from collections import Counter\n\ndef count_words(text):\n    # 统计文本中每个单词的出现次数\n    words = text.lower().split()\n    return Counter(words)",
    "def filter_data(items, condition):\n    # 使用条件函数过滤列表数据\n    return [item for item in items if condition(item)]",

    # 文件操作
    "import json\n\ndef save_json(data, file_path):\n    # 将数据保存为JSON文件\n    with open(file_path, 'w') as f:\n        json.dump(data, f, indent=4)",
    "import csv\n\ndef read_csv_to_dict(file_path):\n    # 读取CSV文件并转换为字典列表\n    result = []\n    with open(file_path, 'r') as f:\n        reader = csv.DictReader(f)\n        for row in reader:\n            result.append(row)\n    return result",

    # Web开发
    "from flask import Flask, jsonify\n\napp = Flask(__name__)\n\n@app.route('/')\ndef home():\n    # 定义一个简单的Flask API端点\n    return jsonify({'message': 'Hello World!'})",
    "import requests\n\ndef fetch_api(url, params=None):\n    # 发送HTTP GET请求并返回JSON响应\n    response = requests.get(url, params=params)\n    if response.status_code == 200:\n        return response.json()\n    else:\n        return None",

    # 机器学习
    "from sklearn.model_selection import train_test_split\n\ndef split_data(X, y, test_size=0.2):\n    # 分割数据集为训练集和测试集\n    return train_test_split(X, y, test_size=test_size, random_state=42)",
    "from sklearn.linear_model import LogisticRegression\n\ndef train_logistic_regression(X_train, y_train):\n    # 训练逻辑回归模型\n    model = LogisticRegression()\n    model.fit(X_train, y_train)\n    return model",

    # 异步编程
    "import asyncio\n\nasync def fetch_data(url):\n    # 异步获取数据\n    await asyncio.sleep(1)\n    return {'data': url}",
    "async def main():\n    # 异步任务主函数\n    tasks = [fetch_data(url) for url in ['a.com', 'b.com']]\n    results = await asyncio.gather(*tasks)\n    return results",

    # 设计模式
    "class Singleton:\n    # 单例模式实现\n    _instance = None\n    def __new__(cls):\n        if not cls._instance:\n            cls._instance = super().__new__(cls)\n        return cls._instance",
    "class Observer:\n    # 观察者模式实现\n    def __init__(self):\n        self.observers = []\n    def add_observer(self, observer):\n        # 添加观察者\n        self.observers.append(observer)\n    def notify(self, event):\n        # 通知所有观察者\n        for observer in self.observers:\n            observer.update(event)",

    # 工具函数
    "def retry(func, max_attempts=3):\n    # 重试装饰器，用于函数执行失败时重试\n    for attempt in range(max_attempts):\n        try:\n            return func()\n        except Exception as e:\n            if attempt == max_attempts - 1:\n                raise e",
    "def memoize(func):\n    # 记忆化装饰器，缓存函数结果\n    cache = {}\n    def wrapper(*args):\n        if args not in cache:\n            cache[args] = func(*args)\n        return cache[args]\n    return wrapper",

    # 数学与算法
    "def binary_search(arr, target):\n    # 二分查找算法\n    left, right = 0, len(arr) - 1\n    while left <= right:\n        mid = (left + right) // 2\n        if arr[mid] == target:\n            return mid\n        elif arr[mid] < target:\n            left = mid + 1\n        else:\n            right = mid - 1\n    return -1",
    "def fibonacci(n):\n    # 斐波那契数列\n    a, b = 0, 1\n    for _ in range(n):\n        a, b = b, a + b\n    return a",

    # 类与面向对象
    "class Rectangle:\n    # 矩形类，计算面积和周长\n    def __init__(self, width, height):\n        self.width = width\n        self.height = height\n    def area(self):\n        return self.width * self.height\n    def perimeter(self):\n        return 2 * (self.width + self.height)",
    "class Animal:\n    # 动物基类\n    def __init__(self, name):\n        self.name = name\n    def speak(self):\n        # 子类必须实现的抽象方法\n        raise NotImplementedError('Subclasses must implement this method')\n\nclass Dog(Animal):\n    # 狗类，继承自动物基类\n    def speak(self):\n        return 'Woof!'"
    # 导入必要模块
    "import pygame\nimport sys\nimport math\nfrom pygame.locals import *\n\n# 初始化pygame\npygame.init()\n\n# 设置中文字体\npygame.font.init()\nfont = pygame.font.SysFont(['SimHei', 'WenQuanYi Micro Hei', 'Heiti TC'], 24)",

    # 游戏初始化
    "def init_game(width, height, title):\n    # 创建游戏窗口并返回屏幕对象和时钟\n    screen = pygame.display.set_mode((width, height))\n    pygame.display.set_caption(title)\n    clock = pygame.time.Clock()\n    return screen, clock",

    # 加载资源
    "def load_image(path, size=None):\n    # 加载图像并根据需要调整大小\n    image = pygame.image.load(path).convert_alpha()\n    if size:\n        image = pygame.transform.scale(image, size)\n    return image",

    # 创建精灵类
    "class Player(pygame.sprite.Sprite):\n    def __init__(self, x, y, speed=5):\n        super().__init__()\n        self.image = pygame.Surface((50, 50))\n        self.image.fill((0, 255, 0))\n        self.rect = self.image.get_rect()\n        self.rect.center = (x, y)\n        self.speed = speed\n        self.health = 100\n        self.direction = 0  # 角度\n\n    def update(self, keys):\n        # 处理玩家移动\n        dx, dy = 0, 0\n        if keys[K_w]:\n            dy -= self.speed\n        if keys[K_s]:\n            dy += self.speed\n        if keys[K_a]:\n            dx -= self.speed\n        if keys[K_d]:\n            dx += self.speed\n        self.rect.x += dx\n        self.rect.y += dy\n\n    def rotate_to_mouse(self):\n        # 使玩家朝向鼠标位置\n        mouse_x, mouse_y = pygame.mouse.get_pos()\n        dx = mouse_x - self.rect.centerx\n        dy = mouse_y - self.rect.centery\n        self.direction = math.degrees(math.atan2(dy, dx))\n\n    def shoot(self, bullet_group):\n        # 发射子弹\n        bullet = Bullet(self.rect.centerx, self.rect.centery, self.direction)\n        bullet_group.add(bullet)",

    # 子弹类
    "class Bullet(pygame.sprite.Sprite):\n    def __init__(self, x, y, direction):\n        super().__init__()\n        self.image = pygame.Surface((10, 10))\n        self.image.fill((255, 0, 0))\n        self.rect = self.image.get_rect()\n        self.rect.center = (x, y)\n        self.speed = 10\n        self.direction = direction\n        # 计算速度分量\n        self.vx = math.cos(math.radians(direction)) * self.speed\n        self.vy = math.sin(math.radians(direction)) * self.speed\n        self.lifetime = 60  # 子弹生命周期\n\n    def update(self):\n        # 更新子弹位置\n        self.rect.x += self.vx\n        self.rect.y += self.vy\n        self.lifetime -= 1\n        # 如果子弹生命周期结束或超出屏幕，删除子弹\n        if self.lifetime <= 0 or not pygame.display.get_surface().get_rect().colliderect(self.rect):\n            self.kill()",

    # 敌人类
    "class Enemy(pygame.sprite.Sprite):\n    def __init__(self, x, y, speed=2):\n        super().__init__()\n        self.image = pygame.Surface((40, 40))\n        self.image.fill((255, 0, 0))\n        self.rect = self.image.get_rect()\n        self.rect.center = (x, y)\n        self.speed = speed\n        self.health = 3\n\n    def update(self, player_pos):\n        # 敌人追踪玩家\n        dx = player_pos[0] - self.rect.centerx\n        dy = player_pos[1] - self.rect.centery\n        distance = max(1, math.sqrt(dx*dx + dy*dy))  # 避免除以零\n        self.rect.x += (dx / distance) * self.speed\n        self.rect.y += (dy / distance) * self.speed\n\n    def take_damage(self, damage=1):\n        # 敌人受伤\n        self.health -= damage\n        if self.health <= 0:\n            self.kill()",

    # 准星类
    "class Crosshair(pygame.sprite.Sprite):\n    def __init__(self):\n        super().__init__()\n        self.image = pygame.Surface((20, 20), pygame.SRCALPHA)\n        pygame.draw.circle(self.image, (255, 255, 255), (10, 10), 10, 2)\n        self.rect = self.image.get_rect()\n\n    def update(self):\n        # 准星跟随鼠标移动\n        self.rect.center = pygame.mouse.get_pos()",

    # 粒子效果类
    "class Particle(pygame.sprite.Sprite):\n    def __init__(self, x, y, color=(255, 255, 0)):\n        super().__init__()\n        size = random.randint(2, 6)\n        self.image = pygame.Surface((size, size), pygame.SRCALPHA)\n        pygame.draw.circle(self.image, color, (size//2, size//2), size//2)\n        self.rect = self.image.get_rect()\n        self.rect.center = (x, y)\n        # 随机速度\n        self.vx = random.uniform(-3, 3)\n        self.vy = random.uniform(-3, 3)\n        self.lifetime = random.randint(10, 30)\n\n    def update(self):\n        # 更新粒子位置和生命周期\n        self.rect.x += self.vx\n        self.rect.y += self.vy\n        self.lifetime -= 1\n        if self.lifetime <= 0:\n            self.kill()",

    # 创建粒子爆炸效果
    "def create_particle_explosion(x, y, count=10, color=(255, 255, 0)):\n    particles = pygame.sprite.Group()\n    for _ in range(count):\n        particle = Particle(x, y, color)\n        particles.add(particle)\n    return particles",

    # 碰撞检测
    "def check_collisions(bullet_group, enemy_group, explosion_group):\n    # 检测子弹与敌人的碰撞\n    collisions = pygame.sprite.groupcollide(bullet_group, enemy_group, True, False)\n    for enemy in collisions.values():\n        for e in enemy:\n            e.take_damage()\n            # 创建爆炸效果\n            explosion = create_particle_explosion(e.rect.centerx, e.rect.centery)\n            explosion_group.add(explosion.sprites())",

    # 显示文字
    "def draw_text(screen, text, x, y, color=(255, 255, 255), font_size=24):\n    # 在屏幕上绘制文字\n    font = pygame.font.SysFont(['SimHei', 'WenQuanYi Micro Hei', 'Heiti TC'], font_size)\n    text_surface = font.render(text, True, color)\n    screen.blit(text_surface, (x, y))",

    # 主游戏循环
    "def game_loop():\n    # 游戏主循环\n    screen, clock = init_game(800, 600, '射击游戏')\n    pygame.mouse.set_visible(False)  # 隐藏鼠标\n\n    # 创建精灵组\n    all_sprites = pygame.sprite.Group()\n    bullets = pygame.sprite.Group()\n    enemies = pygame.sprite.Group()\n    explosions = pygame.sprite.Group()\n\n    # 创建玩家和准星\n    player = Player(400, 300)\n    crosshair = Crosshair()\n    all_sprites.add(player)\n    all_sprites.add(crosshair)\n\n    # 创建敌人计时器\n    enemy_timer = 0\n    ENEMY_SPAWN_RATE = 60  # 帧率\n\n    # 游戏变量\n    score = 0\n    running = True\n\n    while running:\n        # 事件处理\n        for event in pygame.event.get():\n            if event.type == QUIT:\n                running = False\n            elif event.type == MOUSEBUTTONDOWN:\n                player.shoot(bullets)\n                all_sprites.add(bullets)\n\n        # 生成敌人\n        enemy_timer += 1\n        if enemy_timer >= ENEMY_SPAWN_RATE:\n            enemy = Enemy(random.randint(0, 800), random.randint(0, 600))\n            enemies.add(enemy)\n            all_sprites.add(enemy)\n            enemy_timer = 0\n\n        # 获取按键状态\n        keys = pygame.key.get_pressed()\n        player.update(keys)\n        player.rotate_to_mouse()\n        crosshair.update()\n        bullets.update()\n        enemies.update(player.rect.center)\n        explosions.update()\n\n        # 碰撞检测\n        check_collisions(bullets, enemies, explosions)\n\n        # 绘制\n        screen.fill((0, 0, 0))\n        all_sprites.draw(screen)\n        explosions.draw(screen)\n        draw_text(screen, f'分数: {score}', 10, 10)\n        draw_text(screen, f'生命值: {player.health}', 10, 40)\n\n        # 更新显示\n        pygame.display.flip()\n        clock.tick(60)\n\n    pygame.quit()\n    sys.exit()\n\nif __name__ == '__main__':\n    game_loop()"
    # 带缩进提示的训练数据示例
    "def calculate_sum(numbers):\n    # 函数体必须缩进（4空格）\n    total = 0\n    for num in numbers:\n        # 循环体必须缩进（嵌套4空格）\n        total += num\n    return total",
    "if score > 90:\n    # if条件后必须缩进\n    print('优秀')\nelse:\n    # else后必须缩进\n    print('继续努力')",
]

# 初始化模型
ai = CodeGeneratorAI(
    max_sequence_length=1000,  # 控制生成代码的最大长度
    vocab_size=3000          # 词汇表大小，根据数据量调整
)

# 训练模型（50轮就可以看到效果，更多轮次效果更好）
print("开始训练模型...")
ai.train(
    code_samples=training_data,
    epochs=50,
    batch_size=16
)

# 保存模型，这样下次就不用重新训练了
# 保存模型（无警告版本）
print("保存模型...")
ai.model.save('code_model.keras')  # 使用.keras格式

# 保存分词器配置（保持不变）
import json
with open('tokenizer_config.json', 'w') as f:
    f.write(ai.tokenizer.to_json())

print("模型训练并保存完成！")